import logging
import threading

import pika
import yaml


class MqTools:
    def __init__(self, ymls):
        self.convert_task_channel = None
        self.train_task_channel = None
        self.train_task_result_channel = None
        self.model_update_reuslt_channel = None
        self.model_update_channel = None
        self.batch_task_result_channel = None
        self.batch_task_channel = None
        self.queues_names = None
        self.vhost = None
        self.password = None
        self.user = None
        self.port = None
        self.host = None
        self.ymls = yaml.load(open(ymls, 'r'), Loader=yaml.FullLoader) 
        self.read_config()

    def read_config(self):
        self.host = self.ymls['host']
        self.port = self.ymls['port']
        self.user = self.ymls['user']
        self.password = self.ymls['password']
        self.vhost = self.ymls['vhost']
        self.queues_names = self.ymls['queues']

    def make_connection(self):
        
        credentials = pika.PlainCredentials(self.user, self.password)
        parameters = pika.ConnectionParameters(host=self.host, port=self.port, virtual_host=self.vhost, credentials=credentials)
        connection = pika.BlockingConnection(parameters)
        return connection

    def make_channel(self, queue_name):
        connection = self.make_connection()
        self.batch_task_result_channel = connection.channel()
        self.batch_task_result_channel.queue_declare(queue=queue_name['producer']['batch_task_result'], durable=True)

        self.train_task_result_channel = connection.channel()
        self.train_task_result_channel.queue_declare(queue=queue_name['producer']['train_tast_result'], durable=True)

        self.model_update_reuslt_channel = connection.channel()
        self.model_update_reuslt_channel.queue_declare(queue=queue_name['producer']['model_update_reuslt'], durable=True)



    def batch_task_channel_start_consuming(self, queue_name, batch_task_callback):
        self.batch_task_channel = self.make_connection().channel()
        self.batch_task_channel.queue_declare(queue=queue_name['consumer']['batch_task'], durable=True)
        self.batch_task_channel.basic_qos(prefetch_count=1)
        self.batch_task_channel.basic_consume(queue=queue_name['consumer']['batch_task'], on_message_callback=batch_task_callback, auto_ack=False)
        logging.info(' [*] Waiting for messages. To exit press CTRL+C')
        self.batch_task_channel.start_consuming()


    def model_update_channel_start_consuming(self, queue_name, model_update_callback):
        self.model_update_channel = self.make_connection().channel()
        self.model_update_channel.queue_declare(queue=queue_name['consumer']['model_update'], durable=True)
        self.model_update_channel.basic_qos(prefetch_count=1)
        self.model_update_channel.basic_consume(queue=queue_name['consumer']['model_update'], on_message_callback=model_update_callback, auto_ack=False)

    def train_task_channel_start_consuming(self, queue_name, train_task_callback):
        self.train_task_channel = self.make_connection().channel()
        self.train_task_channel.queue_declare(queue=queue_name['consumer']['train_task'], durable=True)
        self.train_task_channel.basic_qos(prefetch_count=1)
        self.train_task_channel.basic_consume(queue=queue_name['consumer']['train_task'], on_message_callback=train_task_callback, auto_ack=False)

    def convert_task_channel_start_consuming(self, queue_name, convert_task_callback):
        self.convert_task_channel = self.make_connection().channel()
        self.convert_task_channel.queue_declare(queue=queue_name['consumer']['convert_task'], durable=True)
        self.convert_task_channel.basic_qos(prefetch_count=1)
        self.convert_task_channel.basic_consume(queue=queue_name['consumer']['convert_task'], on_message_callback=convert_task_callback, auto_ack=False)

    def send(self, queue, message):
        if queue == 'batch_task_result':
            self.batch_task_result_channel.basic_publish(exchange='', routing_key=self.queues_names['producer']['batch_task_result'], body=message)
        elif queue == 'model_update_result':
            self.model_update_reuslt_channel.basic_publish(exchange='', routing_key=self.queues_names['producer']['model_update_result'], body=message)
        elif queue == 'train_task_queue_result':
            self.model_update_reuslt_channel.basic_publish(exchange='', routing_key=self.queues_names['producer']['train_task_queue_result'], body=message)

    def start_consuming(self, batch_task_callback, model_update_callback, train_task_callback, convert_task_callback):
        # self.batch_task_channel_start_consuming(self.queues_names, batch_task_callback)
        # 创建三个线程
        threads = []
        t1 = threading.Thread(target=self.batch_task_channel_start_consuming, args=(self.queues_names, batch_task_callback))
        threads.append(t1)
        t2 = threading.Thread(target=self.train_task_channel_start_consuming(), args=(self.queues_names, train_task_callback))
        threads.append(t2)
        t3 = threading.Thread(target=self.model_update_channel_start_consuming, args=(self.queues_names, model_update_callback))
        threads.append(t3)
        t4 = threading.Thread(target=self.convert_task_channel_start_consuming, args=(self.queues_names, convert_task_callback))
        threads.append(t4)
        # 启动线程
        for t in threads:
            t.setDaemon(True)
            t.start()
        # 等待所有线程完成
        for t in threads:
            t.join()

        try:
            while True:
                continue
        except KeyboardInterrupt:
            logging.info('connection closed')

